
#include "stdafx.h"

#define VLDBUILD         // Declares that we are building Visual Leak Detector.
#include "crtmfcpatch.h" // Provides CRT and MFC patch functions.
#include "vldint.h"      // Provides access to the Visual Leak Detector internals.

#ifndef WIN64
//void * __cdecl operator new(unsigned int,int,char const *,int)
const char    scalar_new_dbg_name[] = "??2@YAPAXIHPBDH@Z";
//void * __cdecl operator new[](unsigned int,int,char const *,int)
const char    vector_new_dbg_name[] = "??_U@YAPAXIHPBDH@Z";
//void * __cdecl operator new(unsigned int)
const char    scalar_new_name[] = "??2@YAPAXI@Z";
//void * __cdecl operator new[](unsigned int)
const char    vector_new_name[] = "??_U@YAPAXI@Z";
#else
//void * __ptr64 __cdecl operator new(unsigned __int64,int,char const * __ptr64,int)
const char    scalar_new_dbg_name[] = "??2@YAPEAX_KHPEBDH@Z";
//void * __ptr64 __cdecl operator new[](unsigned __int64,int,char const * __ptr64,int)
const char    vector_new_dbg_name[] = "??_U@YAPEAX_KHPEBDH@Z";
//void * __ptr64 __cdecl operator new(unsigned __int64)
const char    scalar_new_name[] = "??2@YAPEAX_K@Z";
//void * __ptr64 __cdecl operator new[](unsigned __int64)
const char    vector_new_name[] = "??_U@YAPEAX_K@Z";
#endif

// Global function pointers for explicit dynamic linking with functions listed
// in the import patch table. Using explicit dynamic linking minimizes VLD's
// footprint by loading only modules that are actually used. These pointers will
// be linked to the real functions the first time they are used.

// The import patch table: lists the heap-related API imports that VLD patches
// through to replacement functions provided by VLD. Having this table simply
// makes it more convenient to add additional IAT patches.
patchentry_t VisualLeakDetector::m_kernelbasePatch [] = {
    "GetProcAddress",     NULL, _GetProcAddress, // Not heap related, but can be used to obtain pointers to heap functions.
    NULL,                 NULL, NULL
};

patchentry_t VisualLeakDetector::m_kernel32Patch [] = {
    "HeapAlloc",          NULL, _HeapAlloc,
    "HeapCreate",         NULL, _HeapCreate,
    "HeapDestroy",        NULL, _HeapDestroy,
    "HeapFree",           NULL, _HeapFree,
    "HeapReAlloc",        NULL, _HeapReAlloc,
    NULL,                 NULL, NULL
};

#define ORDINAL(x)          (LPCSTR)x
#if !defined(_M_X64)
#define ORDINAL2(x86, x64)  (LPCSTR)x86
#else
#define ORDINAL2(x86, x64)  (LPCSTR)x64
#endif

VisualLeakDetector::_GetProcAddressType *VisualLeakDetector::m_original_GetProcAddress = NULL;

static patchentry_t mfc42Patch [] = {
    // XXX why are the vector new operators missing for mfc42.dll?
    //ORDINAL(711),         &VC6::pmfcd_scalar_new,              VC6::mfcd_scalar_new,
    NULL,                 NULL,                                 NULL
};

static patchentry_t mfc42dPatch [] = {
    // XXX why are the vector new operators missing for mfc42d.dll?
    ORDINAL(711),         &VC6d::pmfcd_scalar_new,             VC6d::mfcd_scalar_new,
    ORDINAL(712),         &VC6d::pmfcd__scalar_new_dbg_4p,     VC6d::mfcd__scalar_new_dbg_4p,
    ORDINAL(714),         &VC6d::pmfcd__scalar_new_dbg_3p,     VC6d::mfcd__scalar_new_dbg_3p,
    NULL,                 NULL,                                 NULL
};

static patchentry_t mfc42uPatch [] = {
    // XXX why are the vector new operators missing for mfc42u.dll?
    //ORDINAL(711),         &VC6::pmfcud_scalar_new,		        VC6::mfcud_scalar_new,
    NULL,                 NULL,                                 NULL
};

static patchentry_t mfc42udPatch [] = {
    // XXX why are the vector new operators missing for mfc42ud.dll?
    ORDINAL(711),         &VC6d::pmfcud_scalar_new,		    VC6d::mfcud_scalar_new,
    ORDINAL(712),         &VC6d::pmfcud__scalar_new_dbg_4p,    VC6d::mfcud__scalar_new_dbg_4p,
    ORDINAL(714),         &VC6d::pmfcud__scalar_new_dbg_3p,    VC6d::mfcud__scalar_new_dbg_3p,
    NULL,                 NULL,                                 NULL
};

                                                                
static patchentry_t mfc100Patch [] = {						    
    ORDINAL(265),         &VS10::pmfcd_vector_new,		        VS10::mfcd_vector_new,
    ORDINAL2(1294, 1272), &VS10::pmfcd_scalar_new,		        VS10::mfcd_scalar_new,
    NULL,                 NULL,                                 NULL
};
                                                                
static patchentry_t mfc100dPatch [] = {						    
    ORDINAL(267),         &VS10d::pmfcd_vector_new,		    VS10d::mfcd_vector_new,
    ORDINAL(268),         &VS10d::pmfcd__vector_new_dbg_4p,    VS10d::mfcd__vector_new_dbg_4p,
    ORDINAL(269),         &VS10d::pmfcd__vector_new_dbg_3p,    VS10d::mfcd__vector_new_dbg_3p,
    ORDINAL2(1427, 1405), &VS10d::pmfcd_scalar_new,		    VS10d::mfcd_scalar_new,
    ORDINAL2(1428, 1406), &VS10d::pmfcd__scalar_new_dbg_4p,    VS10d::mfcd__scalar_new_dbg_4p,
    ORDINAL2(1429, 1407), &VS10d::pmfcd__scalar_new_dbg_3p,    VS10d::mfcd__scalar_new_dbg_3p,
    NULL,                 NULL,                                 NULL
};
                                                                
static patchentry_t mfc100uPatch [] = {						    
    ORDINAL(265),         &VS10::pmfcud_vector_new,		    VS10::mfcud_vector_new,
    ORDINAL2(1298, 1276), &VS10::pmfcud_scalar_new,		    VS10::mfcud_scalar_new,
    NULL,                 NULL,                                 NULL
};															    
                                                                
static patchentry_t mfc100udPatch [] = {					    
    ORDINAL(267),         &VS10d::pmfcud_vector_new,		    VS10d::mfcud_vector_new,
    ORDINAL(268),         &VS10d::pmfcud__vector_new_dbg_4p,   VS10d::mfcud__vector_new_dbg_4p,
    ORDINAL(269),         &VS10d::pmfcud__vector_new_dbg_3p,   VS10d::mfcud__vector_new_dbg_3p,
    ORDINAL2(1434, 1412), &VS10d::pmfcud_scalar_new,		    VS10d::mfcud_scalar_new,
    ORDINAL2(1435, 1413), &VS10d::pmfcud__scalar_new_dbg_4p,   VS10d::mfcud__scalar_new_dbg_4p,
    ORDINAL2(1436, 1414), &VS10d::pmfcud__scalar_new_dbg_3p,   VS10d::mfcud__scalar_new_dbg_3p,
    NULL,                 NULL,                                 NULL
};
   
static patchentry_t msvcrtPatch [] = {
    scalar_new_dbg_name,  &VC6::pcrtd__scalar_new_dbg,	    VC6::crtd__scalar_new_dbg,
    //vector_new_dbg_name,  &VC6::pcrtd__vector_new_dbg,     VC6::crtd__vector_new_dbg,
    "calloc",             &VC6::pcrtd_calloc,              VC6::crtd_calloc,
    "malloc",             &VC6::pcrtd_malloc,              VC6::crtd_malloc,
    "realloc",            &VC6::pcrtd_realloc,             VC6::crtd_realloc,
    "_strdup",            &VC6::pcrtd__strdup,             VC6::crtd__strdup,
    scalar_new_name,      &VC6::pcrtd_scalar_new,          VC6::crtd_scalar_new,
    //vector_new_name,      &VC6::pcrtd_vector_new,          VC6::crtd_vector_new,
    NULL,                 NULL,                             NULL
};

static patchentry_t msvcrtdPatch [] = {
    "_calloc_dbg",        &VC6d::pcrtd__calloc_dbg,	    VC6d::crtd__calloc_dbg,
    "_malloc_dbg",        &VC6d::pcrtd__malloc_dbg,	    VC6d::crtd__malloc_dbg,
    "_realloc_dbg",       &VC6d::pcrtd__realloc_dbg,       VC6d::crtd__realloc_dbg,
    scalar_new_dbg_name,  &VC6d::pcrtd__scalar_new_dbg,    VC6d::crtd__scalar_new_dbg,
    //vector_new_dbg_name,  &VC6d::pcrtd__vector_new_dbg,     VC6d::crtd__vector_new_dbg,
    "calloc",             &VC6d::pcrtd_calloc,             VC6d::crtd_calloc,
    "malloc",             &VC6d::pcrtd_malloc,             VC6d::crtd_malloc,
    "realloc",            &VC6d::pcrtd_realloc,            VC6d::crtd_realloc,
    "_strdup",            &VC6d::pcrtd__strdup,            VC6d::crtd__strdup,
    scalar_new_name,      &VC6d::pcrtd_scalar_new,         VC6d::crtd_scalar_new,
    //vector_new_name,      &VC6d::pcrtd_vector_new,         VC6d::crtd_vector_new,
    NULL,                 NULL,                             NULL
};

static patchentry_t msvcr100Patch [] = {
    scalar_new_dbg_name,  &VS10::pcrtd__scalar_new_dbg,    VS10::crtd__scalar_new_dbg,
    vector_new_dbg_name,  &VS10::pcrtd__vector_new_dbg,    VS10::crtd__vector_new_dbg,
    "calloc",             &VS10::pcrtd_calloc,             VS10::crtd_calloc,
    "malloc",             &VS10::pcrtd_malloc,             VS10::crtd_malloc,
    "realloc",            &VS10::pcrtd_realloc,            VS10::crtd_realloc,
    "_recalloc",          &VS10::pcrtd_recalloc,           VS10::crtd__recalloc,
    "_strdup",            &VS10::pcrtd__strdup,            VS10::crtd__strdup,
    "_wcsdup",            &VS10::pcrtd__wcsdup,            VS10::crtd__wcsdup,
    scalar_new_name,      &VS10::pcrtd_scalar_new,         VS10::crtd_scalar_new,
    vector_new_name,      &VS10::pcrtd_vector_new,         VS10::crtd_vector_new,
    "_aligned_malloc",			    &VS10::pcrtd_aligned_malloc,               VS10::crtd__aligned_malloc,
    "_aligned_offset_malloc",       &VS10::pcrtd_aligned_offset_malloc,        VS10::crtd__aligned_offset_malloc,
    "_aligned_realloc",             &VS10::pcrtd_aligned_realloc,              VS10::crtd__aligned_realloc,
    "_aligned_offset_realloc",      &VS10::pcrtd_aligned_offset_realloc,       VS10::crtd__aligned_offset_realloc,
    "_aligned_recalloc",            &VS10::pcrtd_aligned_recalloc,             VS10::crtd__aligned_recalloc,
    "_aligned_offset_recalloc",     &VS10::pcrtd_aligned_offset_recalloc,      VS10::crtd__aligned_offset_recalloc,
    NULL,                           NULL,                                       NULL,                                       
};

static patchentry_t msvcr100dPatch [] = {
    "_calloc_dbg",        &VS10d::pcrtd__calloc_dbg,	    VS10d::crtd__calloc_dbg,
    "_malloc_dbg",        &VS10d::pcrtd__malloc_dbg,	    VS10d::crtd__malloc_dbg,
    "_realloc_dbg",       &VS10d::pcrtd__realloc_dbg,      VS10d::crtd__realloc_dbg,
    "_recalloc_dbg",      &VS10d::pcrtd__recalloc_dbg,     VS10d::crtd__recalloc_dbg,
    "_strdup_dbg",        &VS10d::pcrtd__strdup_dbg,       VS10d::crtd__strdup_dbg,
    "_wcsdup_dbg",        &VS10d::pcrtd__wcsdup_dbg,       VS10d::crtd__wcsdup_dbg,
    scalar_new_dbg_name,  &VS10d::pcrtd__scalar_new_dbg,   VS10d::crtd__scalar_new_dbg,
    vector_new_dbg_name,  &VS10d::pcrtd__vector_new_dbg,   VS10d::crtd__vector_new_dbg,
    "calloc",             &VS10d::pcrtd_calloc,            VS10d::crtd_calloc,
    "malloc",             &VS10d::pcrtd_malloc,            VS10d::crtd_malloc,
    "realloc",            &VS10d::pcrtd_realloc,           VS10d::crtd_realloc,
    "_recalloc",          &VS10d::pcrtd_recalloc,          VS10d::crtd__recalloc,
    "_strdup",            &VS10d::pcrtd__strdup,           VS10d::crtd__strdup,
    "_wcsdup",            &VS10d::pcrtd__wcsdup,           VS10d::crtd__wcsdup,
    scalar_new_name,      &VS10d::pcrtd_scalar_new,        VS10d::crtd_scalar_new,
    vector_new_name,      &VS10d::pcrtd_vector_new,        VS10d::crtd_vector_new,
    "_aligned_malloc_dbg",          &VS10d::pcrtd__aligned_malloc_dbg,         VS10d::crtd__aligned_malloc_dbg,
    "_aligned_offset_malloc_dbg",   &VS10d::pcrtd__aligned_offset_malloc_dbg,  VS10d::crtd__aligned_offset_malloc_dbg,
    "_aligned_realloc_dbg",		    &VS10d::pcrtd__aligned_realloc_dbg,        VS10d::crtd__aligned_realloc_dbg,
    "_aligned_offset_realloc_dbg",  &VS10d::pcrtd__aligned_offset_realloc_dbg, VS10d::crtd__aligned_offset_realloc_dbg,
    "_aligned_recalloc_dbg",		&VS10d::pcrtd__aligned_recalloc_dbg,       VS10d::crtd__aligned_recalloc_dbg,
    "_aligned_offset_recalloc_dbg", &VS10d::pcrtd__aligned_offset_recalloc_dbg,VS10d::crtd__aligned_offset_recalloc_dbg,
    "_aligned_malloc",			    &VS10d::pcrtd_aligned_malloc,              VS10d::crtd__aligned_malloc,
    "_aligned_offset_malloc",       &VS10d::pcrtd_aligned_offset_malloc,       VS10d::crtd__aligned_offset_malloc,
    "_aligned_realloc",             &VS10d::pcrtd_aligned_realloc,             VS10d::crtd__aligned_realloc,
    "_aligned_offset_realloc",      &VS10d::pcrtd_aligned_offset_realloc,      VS10d::crtd__aligned_offset_realloc,
    "_aligned_recalloc",            &VS10d::pcrtd_aligned_recalloc,            VS10d::crtd__aligned_recalloc,
    "_aligned_offset_recalloc",     &VS10d::pcrtd_aligned_offset_recalloc,     VS10d::crtd__aligned_offset_recalloc,
    NULL,                           NULL,                                       NULL,                                       
};

patchentry_t VisualLeakDetector::m_ntdllPatch [] = {
    "RtlAllocateHeap",    NULL, VisualLeakDetector::_RtlAllocateHeap,
    "RtlFreeHeap",        NULL, VisualLeakDetector::_RtlFreeHeap,
    "RtlReAllocateHeap",  NULL, VisualLeakDetector::_RtlReAllocateHeap,
    NULL,                 NULL, NULL
};

patchentry_t VisualLeakDetector::m_ole32Patch [] = {
    "CoGetMalloc",        NULL, VisualLeakDetector::_CoGetMalloc,
    "CoTaskMemAlloc",     NULL, VisualLeakDetector::_CoTaskMemAlloc,
    "CoTaskMemRealloc",   NULL, VisualLeakDetector::_CoTaskMemRealloc,
    NULL,                 NULL, NULL
};

moduleentry_t VisualLeakDetector::m_patchTable [] = {
    // Win32 heap APIs.
    "kernel32.dll", 0x0, m_kernelbasePatch, // we patch this record on Win7
    "kernel32.dll", 0x0, m_kernel32Patch,

    // MFC new operators (exported by ordinal).
    "mfc42.dll",    0x0, mfc42Patch,
    "mfc42d.dll",   0x0, mfc42dPatch,
    "mfc42u.dll",   0x0, mfc42uPatch,
    "mfc42ud.dll",  0x0, mfc42udPatch,
    "mfc100.dll",   0x0, mfc100Patch,
    "mfc100d.dll",  0x0, mfc100dPatch,
    "mfc100u.dll",  0x0, mfc100uPatch,
    "mfc100ud.dll", 0x0, mfc100udPatch,

    // CRT new operators and heap APIs.
    "msvcrt.dll",   0x0, msvcrtPatch,
    "msvcrtd.dll",  0x0, msvcrtdPatch,
    "msvcr100.dll", 0x0, msvcr100Patch,
    "msvcr100d.dll",0x0, msvcr100dPatch,

    // NT APIs.
    "ntdll.dll",    0x0, m_ntdllPatch,

    // COM heap APIs.
    "ole32.dll",    0x0, m_ole32Patch
};
